#!/usr/bin/env python
import sys
sys.path.append('./data_parsers')
from surface_parse_bufkit import bufkit_parser_time_height, bufkit_parser
from datetime import datetime, timedelta
import numpy as np
import cPickle

#infile = '/home/disk/pvort/lmadaus/modsnd/sref_em_ctl_ksea.buf'
infile = 'nam_kbil.buf'
UTC_offset = 8

def main():
    plot_timeheight(infile,'SKNT','NAM','KBIL', None)

def get_timelist():
    nowtime = datetime.now()
    stime = nowtime.replace(hour=6,minute=0,second=0,microsecond=0)
    stime = stime + timedelta(hours=24)
    start_ftime = stime
    start_ftime_wind = stime + timedelta(hours=(UTC_offset-6))
    lead_time = timedelta(hours=18)
    fcst_window = timedelta(hours=24)
    end_ftime = start_ftime + fcst_window
    end_ftime_wind = start_ftime_wind + fcst_window

    start_ptime = start_ftime - lead_time
    end_ptime = end_ftime + timedelta(hours=6)
    
    # Make a list of times to plot
    timelist = []
    t = start_ptime
    fcst_delta = timedelta(hours=1)
    while t <= end_ptime:
        timelist.append(t)
        t = t + fcst_delta
    return timelist, start_ftime, end_ftime, start_ftime_wind, end_ftime_wind



def plot_timeheight(infile, var, model, siteid, export_dir=None, savedir=None):
    fcst_dict = bufkit_parser_time_height(infile)

    # Grab the surface pressure trace
    surface_dict = bufkit_parser(infile)

    # Loop here to remove unwanted times
    timelist,start_ftime,end_ftime,start_ftime_wind,end_ftime_wind = get_timelist()

    all_times = fcst_dict.keys()
    all_times.sort()
    model_init = all_times[0]
    starthr = model_init.hour

    # Now clean up times we don't care about
    for key in fcst_dict.keys():
        if key not in timelist:
            del fcst_dict[key]

    # Make the model uppercase
    model = model.upper()

    # Get number of times
    ntimes = len(fcst_dict.keys())

    # Find vertical resolution
    pres_res = fcst_dict[fcst_dict.keys()[0]]['PRES'][:]
    vert_res = abs(int(pres_res[1]) - int(pres_res[0]))
    del pres_res
    # Height will go from 200 to 1050
    plevs = range(200,1050,vert_res)
    plevs.reverse()
    # Make an empty numpy array
    plot_array = np.empty((len(plevs),ntimes))
    plot_array[:] = np.nan

    # Get a sorted list of dates
    date_sorted = fcst_dict.keys()
    date_sorted.sort()
    

    psfc_trace = [surface_dict[d].press for d in date_sorted]
    sfc_uwnd_trace = [surface_dict[d].uwnd for d in date_sorted]
    sfc_vwnd_trace = [surface_dict[d].vwnd for d in date_sorted] 
    # Dump the dictionary here for column model
    cPickle.dump(fcst_dict,open('NAM_fcst_dict.pickle','w'))
    aux_list = {}
    uv_list = {}
    for dtime,idx in zip(date_sorted, range(len(date_sorted))):
        # Load the data
        if var == 'DWP-DEP':
            cur_profile_T = fcst_dict[dtime]['TMPC']
            cur_profile_D = fcst_dict[dtime]['DWPC']
            cur_profile = [a-b for a,b in zip(cur_profile_T,cur_profile_D)]
        else:
            cur_profile = fcst_dict[dtime][var]
        cur_plevs = fcst_dict[dtime]['PRES']
        # Compute some special variables for wind
        if var == 'SKNT':
            cur_hghts = fcst_dict[dtime]['HGHT']
            cur_profile_T = fcst_dict[dtime]['TMPC']
            cur_dirs = fcst_dict[dtime]['DRCT']
            # Need to search for the height at which the
            # lapse rate is no longer close to dry adiabatic
            hghts = np.array(cur_hghts)
            presses = np.array(cur_plevs)
            temps = np.array(cur_profile_T)
            windvels = np.array(cur_profile)
            # Compute the gradients of height and temperature
            dh = np.gradient(hghts)
            dT = np.gradient(temps)
            dh_km = dh / 1000.
            dTdZ  = np.divide(dT,dh_km)
            # Start at beginning, search for place
            # Where dTdZ > -9
            hgt_idx = 0
            found = False
            for gval,ix in zip(dTdZ, range(len(dTdZ))):
                if gval > -6.0 and not found:
                    found = True
                    hgt_idx = ix+1
            aux_list[dtime] = (hghts[hgt_idx],presses[hgt_idx],np.max(windvels[0:hgt_idx+1]),np.mean(windvels[0:hgt_idx+1]))

            cur_uwnd = []
            cur_vwnd = []
            # Get UV wind components at certain heights for barbs
            wind_hghts = [1000., 950., 900., 850., 800., 750.]
            for hght in wind_hghts:
                if hght in cur_plevs:
                    hgtdex = cur_plevs.index(hght)
                    u_wind = windvels[hgtdex] * -1 * np.sin(cur_dirs[hgtdex] * np.pi/180.)
                    v_wind = windvels[hgtdex] * -1 * np.cos(cur_dirs[hgtdex] * np.pi/180.)
                    cur_uwnd.append(u_wind)
                    cur_vwnd.append(v_wind)
                else:
                    cur_uwnd.append(np.nan)
                    cur_vwnd.append(np.nan)
            try:
                uwstack = np.vstack((uwstack,np.array(cur_uwnd)))
                vwstack = np.vstack((vwstack,np.array(cur_vwnd)))
            except:
                uwstack = np.array(cur_uwnd)
                vwstack = np.array(cur_vwnd)

            
        #print dtime, hgt_idx, hghts[hgt_idx], dTdZ
        #print "cur_plevs:", cur_plevs
        #print "plevs:", plevs
        
        # Find the starting index
        vert_start = plevs.index(cur_plevs[0])
  
        #print cur_profile
        # Now paste this array in
        plot_array[vert_start:len(cur_profile)+vert_start,idx] = cur_profile[:]

    #plot_array = plot_array[5:len(cur_profile)-5,:]
    #print np.shape(plot_array)

    # Now search for rows that are all nan
    bad_rows = []
    for r in range(len(plevs)):
        if np.isnan(plot_array[r,:]).all():
            bad_rows.append(r)
        # Another check to get rid of upper levels
        # in temperature plot
        elif var in ['TMPC','SKNT']:
            if float(plevs[r]) < 650.:
                bad_rows.append(r)

    plot_array = np.delete(plot_array,bad_rows,0)
    plevs = np.array(plevs)
    plevs = np.delete(plevs,bad_rows,None)

    #print len(cur_profile)
    
    #print np.shape(plot_array)

    masked_array = np.ma.array(plot_array,mask=np.isnan(plot_array))
    #print masked_array

        


    import matplotlib
    import matplotlib.pyplot as plt
    import matplotlib.cm as cm
    import matplotlib.dates as mdates 
    from mpl_toolkits.axes_grid1 import make_axes_locatable
    matplotlib.rcParams['xtick.minor.pad']=15


    # Set up the color map dictionary
    cmapd = {'CFRL'    : cm.Blues_r,
             'DWP-DEP' : cm.BrBG_r,
             'OMEG'    : cm.RdBu_r}

    # Set up a clim dictionary
    climd = {'DWP-DEP'    : (0.0,10.0),
             'CFRL'       : (0,100),
             'OMEG'       : (-1.5,1.5),
             'TMPC'       : (-20,25)}

    # do a pcolor plot
    fig = plt.figure(figsize=(12,11))
    try:
        cmap = cmapd[var]
    except:
        cmap = cm.jet
    dates = mdates.date2num(date_sorted)


    if var in climd.keys():
        img = plt.pcolormesh(dates,plevs,masked_array,cmap=cmap,vmin=climd[var][0],\
                             vmax=climd[var][1],edgecolors='None')
    else:
        img = plt.pcolormesh(dates,plevs,masked_array,cmap=cmap,edgecolors='None')

    if var == 'TMPC':
        # Contour the zero-degree Celsius line
        plt.hold(True)
        zero_line = plt.contour(dates,plevs,masked_array,[0.0],colors='b',linewidths=3,linestyles='dashed')
        plt.clabel(zero_line,colors='b',inline_spacing=1,fmt='%1.0f C',rightside_up=True)
    elif var == 'SKNT':
        # Plot the top of the boundary layer
        plt.hold(True)
        blhghts = [aux_list[d][1] for d in date_sorted]
        bl_line = plt.plot(dates,blhghts,c='m',linewidth=2,linestyle='dashed')
        # And overlay wind barbs
        datemesh, hgtmesh = np.meshgrid(dates,wind_hghts)
        wind_barbs = plt.barbs(datemesh, hgtmesh, uwstack.T, vwstack.T)
        #print sfc_uwnd_trace
        #surface_barbs = plt.barbs(dates,psfc_trace,sfc_uwnd_trace, sfc_vwnd_trace)

    plt.xlabel('Forecast time')
    plt.ylabel('Pressure Level (mb)')
    plt.title('%s Forecast %s Time-Height at %s' % (model,var,siteid.upper()))

    # Now blot out the "below ground" values
    plt.hold(True)
    ax = plt.gca()
    plt.fill_between(dates, psfc_trace, [ax.get_ylim()[1]]*len(psfc_trace), facecolor='k')

    # Add some vertical lines for reference
    if var in ['SKNT','WDIR']:
        plt.axvline(x=mdates.date2num(start_ftime_wind),linewidth=2,color='k')
        plt.axvline(x=mdates.date2num(end_ftime_wind),linewidth=2,color='k')
    else:
        plt.axvline(x=mdates.date2num(start_ftime),linewidth=2,color='k')
        plt.axvline(x=mdates.date2num(end_ftime),linewidth=2,color='k')

    ax = plt.gca()
    ax.set_ylim((min(plevs),max(plevs)))
    ax.invert_yaxis()
    ax.set_xlim((min(dates),max(dates)))
    ax.xaxis.set_major_locator(mdates.HourLocator(range(0,25,3)))
    ax.xaxis.set_minor_locator(mdates.DayLocator())
    ax.xaxis.set_major_formatter(mdates.DateFormatter('%HZ'))
    ax.xaxis.set_minor_formatter(matplotlib.dates.DateFormatter('%h %d'))
    plt.grid()
    if var == 'SKNT':
        plt.legend([bl_line[0]],['Mixed Layer Height'],loc=2)
        # Dump the archive of wind values for later ts plot
        top_wind = [aux_list[d][2] for d in date_sorted]
        mean_wind = [aux_list[d][3] for d in date_sorted]
        try:
            if savedir != None:
                f = open('%s/%s_bl_winds.pickle' % (savedir,siteid.upper()),'r')
            else:
                f = open('%s_bl_winds.pickle' % (siteid.upper()),'r')
            windd = cPickle.load(f)
            f.close()
        except:
            windd = {}

        # Dump it as this model

        print "Dumping ",model," at %dZ" % starthr 
        windd['%s %02dZ' % (model, starthr)] = {'dates' : date_sorted,
                                              'top_wind' : top_wind,
                                              'mean_wind' : mean_wind,
                                              'model_init': model_init}
        if savedir != None:
            f = open('%s/%s_bl_winds.pickle' % (savedir,siteid.upper()),'w')
        else:
            f = open('%s_bl_winds.pickle' % siteid.upper(),'w')
        cPickle.dump(windd,f)
        f.close()
        
        
    #ax.format_xdata = mdates.DateFormatter('%m/%d - %HZ')
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right",size="3%",pad=0.1)
    if var == 'DWP-DEP':
        plt.colorbar(img, cax=cax, extend='max') 

    else:
        plt.colorbar(img, cax=cax)


    fig.autofmt_xdate()
    if export_dir != None:
        import os
        plt.savefig('%s_%s_%s_tmht.png' % (siteid.upper(),model.upper(), \
                     var.upper()), bbox_inches='tight')
        os.system('mv %s_%s_%s_tmht.png %s' % (siteid.upper(),model.upper(),\
                     var.upper(), export_dir))
    else:
        plt.show()



if __name__ == '__main__':
    main()
